#!/usr/bin/env python
# -*- coding: ASCII -*-

"""

:Author: Martin Kircher
:Contact: Martin.Kircher@eva.mpg.de
:Date: *20.04.2009

"""
import sys,gzip
import os
from optparse import OptionParser, OptionGroup

indexes = {}
names = {}
index_length = None
ireadlength = None

next_bases = 'ATCTCGTATGCCGTCTTCTGCTTG'
nucleotides = ["A", "C", "G", "T"]


def read_index_file(filename, skipFirst=True, oneDist=True, nbases=True, flength=None):
    global next_bases
    global nucleotides
    global options
    infile = open(filename)
    mlindex = None
    for line in infile:
        if len(line.strip()) > 0:
            fields = line.split()
            if len(fields) >= 2:
                cindex = fields[0].upper()
                # SKIP NON VALID LINES
                valid = True
                for elem in cindex:
                    if elem not in nucleotides:
                        valid = False
                        break
                if not valid: continue
                if (mlindex == None) or (len(cindex) > mlindex): mlindex = len(cindex)
    if (flength > 0) and (flength > mlindex): mlindex = flength

    res = {}
    names = {}
    names[None] = "conflict"
    infile = open(filename)
    for line in infile:
        if len(line.strip()) > 0:
            fields = line.split()
            if len(fields) >= 2:
                cindex = fields[0].upper()
                # SKIP NON VALID LINES
                valid = True
                for elem in cindex:
                    if elem not in nucleotides:
                        valid = False
                        break
                if not valid: continue

                if mlindex != len(cindex):
                    cindex += next_bases[:(mlindex - len(cindex))]
                names[cindex] = fields[1].split("-")[-1]
                if (cindex in res) and (res[cindex] != cindex):
                    print "Error: Index causes a conflict %s (%s) with variant of %s (%s)" % (
                        cindex, names[cindex], res[cur], names[res[cur]])
                    res[cindex] = None
                else:
                    res[cindex] = cindex

                # SKIP FIRST BASE
                if skipFirst:
                    cur = cindex[1:] + next_bases[(mlindex - len(fields[0]))]
                    # print fields[0].upper()+'\t'+cindex+'\t'+cur
                    if (cur in res) and (res[cur] != cindex):
                        print "Error: Skipping first base in index causes a conflict %s (%s)->%s with variant of %s (%s)" % (
                            cindex, names[cindex], cur, res[cur], names[res[cur]])
                        res[cur] = None
                        # sys.exit()
                    else:
                        res[cur] = cindex
                # CREATE 1nt MUTANTS
                if oneDist:
                    for cpos in range(len(cindex)):
                        for base in nucleotides:
                            cur = list(cindex)
                            cur[cpos] = base
                            cur = "".join(cur)
                            if (cur in res) and (res[cur] != cindex):
                                print "Error: Creating mutants of indexes causes a conflict %s (%s)->%s with variant of %s (%s)" % (
                                    cindex, names[cindex], cur, res[cur], names[res[cur]])
                                res[cur] = None
                                # sys.exit()
                            else:
                                res[cur] = cindex
                # CREATE Ns
                if nbases:
                    for cpos in range(len(cindex)):
                        cur = list(cindex)
                        cur[cpos] = "N"
                        cur = "".join(cur)
                        if (cur in res) and (res[cur] != cindex):
                            print "Error: creating mutants of indexes causes a conflict %s (%s)->%s with variant of %s (%s)" % (
                                cindex, names[cindex], cur, res[cur], names[res[cur]])
                            res[cur] = None
                            sys.exit(2)
                        else:
                            res[cur] = cindex

                # REMOVE SOME BASES AND REPLACE BY N
                if (options.setN != None) and len(options.setN) == len(cindex):
                    cur = list(cindex)
                    nvariants = []
                    for cpos, base in enumerate(options.setN.upper()):
                        if base == "N":
                            cur[cpos] = "N"
                    cur = "".join(cur)
                    if (cur in res) and (res[cur] != cindex):
                        print "Error: creating setN mutant of indexes causes a conflict %s (%s)->%s with variant of %s (%s)" % (
                            cindex, names[cindex], cur, res[cur], names[res[cur]])
                        res[cur] = None
                        sys.exit(2)
                    else:
                        print cur, cindex
                        res[cur] = cindex
                    ncur = cur

                    # CREATE EDITS FOR WITH Ns REPLACED INDEX SEQUENCE
                    # SKIP FIRST BASE
                    if skipFirst:
                        cur = ncur[1:] + next_bases[(mlindex - len(fields[0]))]
                        # print fields[0].upper()+'\t'+cindex+'\t'+cur
                        if (cur in res) and (res[cur] != cindex):
                            print "Error: Skipping first base in index causes a conflict %s (%s)->%s with variant of %s (%s)" % (
                                cindex, names[cindex], cur, res[cur], names[res[cur]])
                            res[cur] = None
                            # sys.exit()
                        else:
                            res[cur] = cindex
                    # CREATE 1nt MUTANTS
                    if oneDist:
                        for cpos in range(len(cindex)):
                            for base in nucleotides:
                                cur = list(ncur)
                                cur[cpos] = base
                                cur = "".join(cur)
                                if (cur in res) and (res[cur] != cindex):
                                    print "Error: Creating mutants of indexes causes a conflict %s (%s)->%s with variant of %s (%s)" % (
                                        cindex, names[cindex], cur, res[cur], names[res[cur]])
                                    res[cur] = None
                                    # sys.exit()
                                else:
                                    res[cur] = cindex
                    # CREATE Ns
                    if nbases:
                        for cpos in range(len(cindex)):
                            cur = list(ncur)
                            cur[cpos] = "N"
                            cur = "".join(cur)
                            if (cur in res) and (res[cur] != cindex):
                                print "Error: creating mutants of indexes causes a conflict %s (%s)->%s with variant of %s (%s)" % (
                                    cindex, names[cindex], cur, res[cur], names[res[cur]])
                                res[cur] = None
                                sys.exit()
                            else:
                                res[cur] = cindex

            else:
                print "Error: Line is not valid: %s" % line.strip()
                sys.exit(2)
    return mlindex, names, res


def read_sequence_file(filename, sec_read_start=None):
    global ireadlength
    infile = gzip.open(filename)
    seqid = None
    seq = ""
    qual = -1
    fastq = None
    for line in infile:
        # print line.strip()
        if (line.startswith("@") and (fastq != False) and ((qual == -1) or (len(qual) == len(seq)))) or (
                line.startswith(">") and (fastq != True)):
            # FIRST LINE DECIDES FILE FORMAT!
            if fastq is None:
                if line.startswith("@"):
                    fastq = True
                else:
                    fastq = False
            if len(seq) > 0:
                if sec_read_start is None:
                    if qual is not None:
                        yield seqid, seq[-ireadlength:], seq[:-ireadlength], qual[:-ireadlength], None, None
                    else:
                        yield seqid, seq[-ireadlength:], seq[:-ireadlength], None, None, None
                else:
                    if qual is not None:
                        yield seqid, seq[(sec_read_start - ireadlength):sec_read_start], seq[:(
                                sec_read_start - ireadlength)], qual[:(sec_read_start - ireadlength)], seq[
                                                                                                       sec_read_start:], qual[
                                                                                                                         sec_read_start:]
                    else:
                        yield seqid, seq[(sec_read_start - ireadlength):sec_read_start], seq[:(
                                sec_read_start - ireadlength)], None, seq[sec_read_start:], None
            seqid = line[1:].strip()
            seq = ""
            qual = None
        elif line.startswith("+") and (qual is None):
            qual = ""
        else:
            if qual is None:
                seq += line.strip()
            else:
                try:
                    qual += line.strip()
                except TypeError:
                    print line
                    exit(1)
    infile.close()
    if len(seq) > 0:
        if sec_read_start == None:
            if qual != None:
                yield seqid, seq[-ireadlength:], seq[:-ireadlength], qual[:-ireadlength], None, None
            else:
                yield seqid, seq[-ireadlength:], seq[:-ireadlength], None, None, None
        else:
            if qual != None:
                yield seqid, seq[(sec_read_start - ireadlength):sec_read_start], seq[
                                                                                 :(sec_read_start - ireadlength)], qual[
                                                                                                                   :(
                                                                                                                           sec_read_start - ireadlength)], seq[
                                                                                                                                                           sec_read_start:], qual[
                                                                                                                                                                             sec_read_start:]
            else:
                yield seqid, seq[(sec_read_start - ireadlength):sec_read_start], seq[:(
                        sec_read_start - ireadlength)], None, seq[sec_read_start:], None
    raise StopIteration


parser = OptionParser("%prog [options] seq_files")
parser.add_option("-i", "--index", dest="index", help="File describing index sequences used")
parser.add_option("-l", "--ireadlength", dest="ireadlength",
                  help="Length of index read (in case it is longer than the indexes in the index file provided)",
                  type="int", default=0)
parser.add_option("--bases_after_index", dest="nextbases",
                  help="Nucleotides in adapter downstream of index (default %s)" % next_bases, default=next_bases)
parser.add_option("-s", "--start", dest="start", help="First base of the second read (default None)", type='int')
parser.add_option("--sep_reads", dest="sep_reads", help="Separate first and second read in two files", default=False,
                  action="store_true")
parser.add_option("-o", "--outdir", dest="outdir", help="Create output files in another directory.")
parser.add_option("-v", "--verbose", dest="verbose", help="Turn all messages on", default=False, action="store_true")
parser.add_option("--summary", dest="summary", help="Show summary for tags identified", default=False,
                  action="store_true")
parser.add_option("--setN", dest="setN",
                  help="Set bases from index read to N (e.g. NyyyyyN removes first and last base)")

group = OptionGroup(parser, "Sequencing error correction")
group.add_option("--no_skip_first_base", dest="skip_first_base", help="Do not consider skipping first base of index",
                 default=True, action="store_false")
group.add_option("--no_mutants", dest="mutants", help="Do not consider index mutants within one base distance",
                 default=True, action="store_false")
group.add_option("--no_Ns", dest="Nbases", help="Do not allow one N in index sequences", default=True,
                 action="store_false")
parser.add_option_group(group)
(options, args) = parser.parse_args()

if ((options.start != None) and (options.start > 0)):
    options.start -= 1
else:
    options.start = None
    options.sep_reads = False

next_bases = options.nextbases.upper()

names, indexes = {}, {}
if (options.index != None) and os.path.isfile(options.index):
    if options.verbose: print "Calculating index variants..."
    index_length, names, indexes = read_index_file(options.index, skipFirst=options.skip_first_base,
                                                   oneDist=options.mutants, nbases=options.Nbases,
                                                   flength=options.ireadlength)
    if options.verbose: print "Created %d putative index variants for %d index sequences." % (
        len(indexes), len(names) - 1)
    names["unknown"] = "unknown"
    ireadlength = max(options.ireadlength, index_length)
elif (options.ireadlength > 0) and ((options.index == None) or (options.index.strip() == "")):
    ireadlength = options.ireadlength
else:
    print "Need valid index file."
    sys.exit(2)

files = args
for compfilename in files:
    if os.path.isfile(compfilename):
        # SEPARATE FILENAME AND PATH
        filename = compfilename.strip().split("/")[-1]
        cpath = "/".join(compfilename.strip().split("/")[:-1])
        if len(cpath) > 0: cpath += "/"
        # CREATE OUTPUT FILES DEPENDING ON CURRENT FILNAME AND OUTPUT PATH SITUATION
        outfiles = {}
        if options.verbose: print "Creating output files..."
        if options.sep_reads:
            if filename.startswith("s_") and filename.endswith("_sequence.txt"):
                clane = filename.split("_")[1]
                if len(names) == 0:
                    outfilename = cpath + "s_%s_1_sequence_noindex.txt" % clane
                    outfilename2 = cpath + "s_%s_2_sequence_noindex.txt" % clane
                    if ((options.outdir <> None) and os.path.isdir(options.outdir)):
                        outfilename = options.outdir + "/" + "s_%s_1_sequence.txt" % clane
                        outfilename2 = options.outdir + "/" + "s_%s_2_sequence.txt" % clane
                    if options.verbose: print "Creating:", outfilename, outfilename2
                    outfiles[None] = [open(outfilename, 'w'), open(outfilename2, 'w'), 0, outfilename, outfilename2]
                else:
                    for tag, name in names.iteritems():
                        outfilename = cpath + "s_%s_1_%s_sequence.txt" % (clane, name)
                        outfilename2 = cpath + "s_%s_2_%s_sequence.txt" % (clane, name)
                        if ((options.outdir <> None) and os.path.isdir(options.outdir)):
                            outfilename = options.outdir + "/" + "s_%s_1_%s_sequence.txt" % (clane, name)
                            outfilename2 = options.outdir + "/" + "s_%s_2_%s_sequence.txt" % (clane, name)
                        if options.verbose: print "Creating:", outfilename, outfilename2
                        outfiles[tag] = [open(outfilename, 'w'), open(outfilename2, 'w'), 0, outfilename, outfilename2]
            else:
                if len(names) == 0:
                    outfilename = cpath + ".".join(filename.split(".")[:-1]) + "_r1.txt"
                    outfilename2 = cpath + ".".join(filename.split(".")[:-1]) + "_r2.txt"
                    if ((options.outdir <> None) and os.path.isdir(options.outdir)):
                        outfilename = options.outdir + "/" + ".".join(filename.split(".")[:-1]) + "_r1.txt"
                        outfilename2 = options.outdir + "/" + ".".join(filename.split(".")[:-1]) + "_r2.txt"
                    if options.verbose: print "Creating:", outfilename, outfilename2
                    outfiles[None] = [open(outfilename, 'w'), open(outfilename2, 'w'), 0, outfilename, outfilename2]
                else:
                    for tag, name in names.iteritems():
                        outfilename = cpath + ".".join(filename.split(".")[:-1]) + "_%s_r1.txt" % name
                        outfilename2 = cpath + ".".join(filename.split(".")[:-1]) + "_%s_r2.txt" % name
                        if ((options.outdir <> None) and os.path.isdir(options.outdir)):
                            outfilename = options.outdir + "/" + ".".join(
                                filename.split(".")[:-1]) + "_%s_r1.txt" % name
                            outfilename2 = options.outdir + "/" + ".".join(
                                filename.split(".")[:-1]) + "_%s_r2.txt" % name
                        if options.verbose: print "Creating:", outfilename, outfilename2
                        outfiles[tag] = [open(outfilename, 'w'), open(outfilename2, 'w'), 0, outfilename, outfilename2]
        else:
            if filename.startswith("s_") and filename.endswith("_sequence.txt"):
                clane = filename.split("_")[1]
                if len(names) == 0:
                    outfilename = cpath + "s_%s_sequence_noindex.txt" % (clane)
                    if ((options.outdir <> None) and os.path.isdir(options.outdir)):
                        outfilename = options.outdir + "/" + "s_%s_sequence.txt" % (clane)
                    if options.verbose: print "Creating:", outfilename
                    outfiles[None] = [open(outfilename, 'w'), 0, outfilename]
                else:
                    for tag, name in names.iteritems():
                        outfilename = cpath + "s_%s_%s_sequence.txt" % (clane, name)
                        if ((options.outdir <> None) and os.path.isdir(options.outdir)):
                            outfilename = options.outdir + "/" + "s_%s_%s_sequence.txt" % (clane, name)
                        if options.verbose: print "Creating:", outfilename
                        outfiles[tag] = [open(outfilename, 'w'), 0, outfilename]
            else:
                if len(names) == 0:
                    outfilename = cpath + ".".join(filename.split(".")[:-1]) + "_noindex.txt"
                    if ((options.outdir <> None) and os.path.isdir(options.outdir)):
                        outfilename = options.outdir + "/" + filename
                    if options.verbose: print "Creating:", outfilename
                    outfiles[None] = [open(outfilename, 'w'), 0, outfilename]
                else:
                    for tag, name in names.iteritems():
                        outfilename = cpath + ".".join(filename.split(".")[:-1]) + "_%s.txt" % name
                        if ((options.outdir <> None) and os.path.isdir(options.outdir)):
                            outfilename = options.outdir + "/" + ".".join(filename.split(".")[:-1]) + "_%s.txt" % name
                        if options.verbose: print "Creating:", outfilename
                        outfiles[tag] = [open(outfilename, 'w'), 0, outfilename]

        # ACTUAL INDEX IDENTIFICATION AND READ SORTING
        for seqid, indseq, r1, q1, r2, q2 in read_sequence_file(compfilename, sec_read_start=options.start):
            indseq = indseq[:ireadlength]
            if ((len(indseq) == 0) or (len(outfiles) == 1)) and (None in outfiles):
                cind = None
            elif indseq in indexes:
                cind = indexes[indseq]
                seqid += "#" + indseq
            else:
                cind = "unknown"
                seqid += "#" + indseq
            if options.sep_reads:
                outfiles[cind][2] += 1
                if q1 != None:
                    outfiles[cind][0].write("@%s/1\n%s\n+%s/1\n%s\n" % (seqid, r1, seqid, q1))
                    outfiles[cind][1].write("@%s/2\n%s\n+%s/2\n%s\n" % (seqid, r2, seqid, q2))
                else:
                    outfiles[cind][0].write(">%s/1\n%s\n" % (seqid, r1))
                    outfiles[cind][1].write(">%s/2\n%s\n" % (seqid, r2))
            else:
                outfiles[cind][1] += 1
                if q1 != None:
                    if r2 != None:
                        outfiles[cind][0].write("@%s\n%s\n+%s\n%s\n" % (seqid, r1 + r2, seqid, q1 + q2))
                    else:
                        outfiles[cind][0].write("@%s\n%s\n+%s\n%s\n" % (seqid, r1, seqid, q1))
                else:
                    if r2 != None:
                        outfiles[cind][0].write(">%s\n%s\n" % (seqid, r1 + r2))
                    else:
                        outfiles[cind][0].write(">%s\n%s\n" % (seqid, r1))

        # CREATE SUMMARY AND CLEAN UP NOT USED FILES
        summary = []
        if options.verbose: print "Cleaning up not needed output files..."
        if options.sep_reads:
            for tag, value in outfiles.iteritems():
                value[0].close()
                value[1].close()
                if value[2] == 0:
                    try:
                        if options.verbose: print "Removing: %s %s" % (value[3], value[4])
                        os.remove(value[3])
                        os.remove(value[4])
                    except:
                        print "Error: Removing output files %s %s"(value[3], value[4])
                elif options.summary:
                    if tag in names:
                        summary.append("  Observed %10d tags for %s (%s)" % (value[2], names[tag], tag))
                    else:
                        summary.append("  Created %10d sequences with index read removed" % value[2])
        else:
            for tag, value in outfiles.iteritems():
                value[0].close()
                if value[1] == 0:
                    try:
                        if options.verbose: print "Removing: %s" % (value[2])
                        os.remove(value[2])
                    except:
                        print "Error: Removing output files %s %s"(value[2])
                elif options.summary:
                    if tag in names:
                        summary.append("  Observed %10d tags for %s (%s)" % (value[1], names[tag], tag))
                    else:
                        summary.append("  Created %10d sequences with index read removed" % value[1])
        summary.sort()
        summary.reverse()
        if options.summary: print "Summary:"
        for line in summary:
            print line
